import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader

import torchvision
import torchvision.transforms as transforms

import matplotlib.pyplot as plt
import numpy as np
from tqdm import tqdm

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

if device.type == 'cuda':
    print(f"   GPU: {torch.cuda.get_device_name(0)}")
    print(f"   Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
else:
    print("   No GPU detected! Training will be slower...")

Using device: cuda
   GPU: Tesla T4
   Memory: 15.6 GB

BATCH_SIZE = 128
LATENT_DIM = 128
LEARNING_RATE = 1e-3
EPOCHS = 50

IMG_CHANNELS = 3
IMG_SIZE = 32

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

train_dataset = torchvision.datasets.CIFAR10(
    root='./data', 
    train=True,
    download=True, 
    transform=transform
)

test_dataset = torchvision.datasets.CIFAR10(
    root='./data', 
    train=False,
    download=True, 
    transform=transform
)

train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)

classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

print(f"Training samples: {len(train_dataset)}")
print(f"Test samples: {len(test_dataset)}")
print(f"Batches per epoch: {len(train_loader)}")

100%|██████████| 170M/170M [00:05<00:00, 31.1MB/s]

Training samples: 50000
Test samples: 10000
Batches per epoch: 391

class Encoder(nn.Module):
    def __init__(self, latent_dim=128):
        super(Encoder, self).__init__()
        
        self.conv_layers = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(32),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(256),
            nn.LeakyReLU(0.2, inplace=True),
        )
        
        self.flatten_size = 256 * 2 * 2
        self.fc_mu = nn.Linear(self.flatten_size, latent_dim)
        self.fc_logvar = nn.Linear(self.flatten_size, latent_dim)
    
    def forward(self, x):
        x = self.conv_layers(x)
        x = x.view(x.size(0), -1)
        mu = self.fc_mu(x)
        logvar = self.fc_logvar(x)
        return mu, logvar

print("Encoder defined")

Encoder defined

class Decoder(nn.Module):
    def __init__(self, latent_dim=128):
        super(Decoder, self).__init__()
        
        self.latent_dim = latent_dim
        
        self.fc = nn.Sequential(
            nn.Linear(latent_dim, 256 * 2 * 2),
            nn.ReLU(inplace=True)
        )
        
        self.deconv_layers = nn.Sequential(
            nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            
            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            
            nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            
            nn.ConvTranspose2d(32, 3, kernel_size=4, stride=2, padding=1),
            nn.Tanh()
        )
    
    def forward(self, z):
        x = self.fc(z)
        x = x.view(x.size(0), 256, 2, 2)
        x = self.deconv_layers(x)
        return x

print("Decoder defined")

Decoder defined

class VAE(nn.Module):
    def __init__(self, latent_dim=128):
        super(VAE, self).__init__()
        
        self.latent_dim = latent_dim
        self.encoder = Encoder(latent_dim)
        self.decoder = Decoder(latent_dim)
    
    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        epsilon = torch.randn_like(std)
        z = mu + std * epsilon
        return z
    
    def forward(self, x):
        mu, logvar = self.encoder(x)
        z = self.reparameterize(mu, logvar)
        x_reconstructed = self.decoder(z)
        return x_reconstructed, mu, logvar
    
    def generate(self, num_samples, device):
        z = torch.randn(num_samples, self.latent_dim).to(device)
        with torch.no_grad():
            samples = self.decoder(z)
        return samples

model = VAE(LATENT_DIM).to(device)

total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"VAE model created!")
print(f"   Total parameters: {total_params:,}")
print(f"   Trainable parameters: {trainable_params:,}")

VAE model created!
   Total parameters: 1,775,939
   Trainable parameters: 1,775,939

def vae_loss(recon_x, x, mu, logvar, beta=1.0):
    batch_size = x.size(0)
    
    recon_loss = F.mse_loss(recon_x, x, reduction='sum') / batch_size
    kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp()) / batch_size
    total_loss = recon_loss + beta * kl_loss
    
    return total_loss, recon_loss, kl_loss

print("Loss function defined")

Loss function defined

def train_vae(model, train_loader, epochs, beta=1.0, lr=1e-3):
    optimizer = optim.Adam(model.parameters(), lr=lr)
    
    history = {
        'total_loss': [],
        'recon_loss': [],
        'kl_loss': []
    }
    
    model.train()
    
    for epoch in range(epochs):
        epoch_total = 0
        epoch_recon = 0
        epoch_kl = 0
        num_batches = 0
        
        pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}")
        
        for batch_idx, (data, _) in enumerate(pbar):
            data = data.to(device)
            
            optimizer.zero_grad()
            recon_batch, mu, logvar = model(data)
            loss, recon_loss, kl_loss = vae_loss(recon_batch, data, mu, logvar, beta)
            loss.backward()
            optimizer.step()
            
            epoch_total += loss.item()
            epoch_recon += recon_loss.item()
            epoch_kl += kl_loss.item()
            num_batches += 1
            
            pbar.set_postfix({
                'loss': f'{loss.item():.4f}',
                'recon': f'{recon_loss.item():.4f}',
                'kl': f'{kl_loss.item():.4f}'
            })
        
        avg_total = epoch_total / num_batches
        avg_recon = epoch_recon / num_batches
        avg_kl = epoch_kl / num_batches
        
        history['total_loss'].append(avg_total)
        history['recon_loss'].append(avg_recon)
        history['kl_loss'].append(avg_kl)
        
        print(f"  → Avg Loss: {avg_total:.4f} | Recon: {avg_recon:.4f} | KL: {avg_kl:.4f}")
    
    return history

print("Training function defined")

Training function defined

print("=" * 60)
print("TRAINING STANDARD VAE (β = 1)")
print("=" * 60)

model = VAE(LATENT_DIM).to(device)
history = train_vae(model, train_loader, EPOCHS, beta=1.0, lr=LEARNING_RATE)

print("\nTraining complete!")

============================================================
TRAINING STANDARD VAE (β = 1)
============================================================

Epoch 1/50:   0%|          | 0/391 [00:00<?, ?it/s]

Epoch 1/50: 100%|██████████| 391/391 [00:15<00:00, 25.78it/s, loss=254.9640, recon=214.2250, kl=40.7390]

  → Avg Loss: 381.9684 | Recon: 346.1202 | KL: 35.8482

Epoch 2/50: 100%|██████████| 391/391 [00:13<00:00, 28.65it/s, loss=221.4289, recon=179.6664, kl=41.7625]

  → Avg Loss: 251.6903 | Recon: 208.1166 | KL: 43.5737

Epoch 3/50: 100%|██████████| 391/391 [00:13<00:00, 28.30it/s, loss=230.7552, recon=181.5909, kl=49.1642]

  → Avg Loss: 231.4502 | Recon: 185.6165 | KL: 45.8337

Epoch 4/50: 100%|██████████| 391/391 [00:13<00:00, 27.98it/s, loss=210.9300, recon=162.3524, kl=48.5777]

  → Avg Loss: 222.4549 | Recon: 174.7670 | KL: 47.6879

def plot_training_curves(history, title="Training Progress"):
    """Plot the training loss curves"""
    fig, axes = plt.subplots(1, 3, figsize=(15, 4))
    
    epochs = range(1, len(history['total_loss']) + 1)
    
    # Total Loss
    axes[0].plot(epochs, history['total_loss'], 'b-', linewidth=2, label='Total Loss')
    axes[0].set_xlabel('Epoch')
    axes[0].set_ylabel('Loss')
    axes[0].set_title('Total Loss (ELBO)')
    axes[0].grid(True, alpha=0.3)
    axes[0].legend()
    
    # Reconstruction Loss
    axes[1].plot(epochs, history['recon_loss'], 'g-', linewidth=2, label='Recon Loss')
    axes[1].set_xlabel('Epoch')
    axes[1].set_ylabel('Loss')
    axes[1].set_title('Reconstruction Loss (MSE)')
    axes[1].grid(True, alpha=0.3)
    axes[1].legend()
    
    # KL Divergence
    axes[2].plot(epochs, history['kl_loss'], 'r-', linewidth=2, label='KL Loss')
    axes[2].set_xlabel('Epoch')
    axes[2].set_ylabel('Loss')
    axes[2].set_title('KL Divergence')
    axes[2].grid(True, alpha=0.3)
    axes[2].legend()
    
    plt.suptitle(title, fontsize=14, fontweight='bold')
    plt.tight_layout()
    plt.savefig('training_curves.png', dpi=150, bbox_inches='tight')
    plt.show()
    
    print("Training curves saved to 'training_curves.png'")

# Plot the training curves
plot_training_curves(history, "Standard VAE (β=1) Training Progress")

Training curves saved to 'training_curves.png'

def generate_image_grid(model, num_images=16, nrow=4, title="Generated Images"):
    """Generate a grid of images from random latent vectors"""
    model.eval()
    
    with torch.no_grad():
        # Sample random latent vectors from standard normal
        z = torch.randn(num_images, model.latent_dim).to(device)
        
        # Generate images
        generated = model.decoder(z)
        
        # Denormalize from [-1, 1] to [0, 1]
        generated = generated * 0.5 + 0.5
        generated = generated.cpu()
    
    # Create grid
    fig, axes = plt.subplots(nrow, nrow, figsize=(10, 10))
    
    for i, ax in enumerate(axes.flat):
        if i < num_images:
            img = generated[i].numpy().transpose(1, 2, 0)
            img = np.clip(img, 0, 1)
            ax.imshow(img)
        ax.axis('off')
    
    plt.suptitle(title, fontsize=16, fontweight='bold')
    plt.tight_layout()
    plt.savefig('generated_grid.png', dpi=150, bbox_inches='tight')
    plt.show()
    
    print("Generated image grid saved to 'generated_grid.png'")

# Generate 16 random images in a 4x4 grid
print("Generating 16 images from random noise...")
generate_image_grid(model, num_images=16, nrow=4, title="VAE Generated Images (β=1)")

Generating 16 images from random noise...

Generated image grid saved to 'generated_grid.png'

def show_reconstructions(model, data_loader, num_images=8):
    """Show original images vs their reconstructions"""
    model.eval()
    
    # Get a batch of test images
    images, labels = next(iter(data_loader))
    images = images[:num_images].to(device)
    
    with torch.no_grad():
        recon, _, _ = model(images)
    
    # Denormalize
    images = images * 0.5 + 0.5
    recon = recon * 0.5 + 0.5
    
    # Plot
    fig, axes = plt.subplots(2, num_images, figsize=(14, 4))
    
    for i in range(num_images):
        # Original
        orig = images[i].cpu().numpy().transpose(1, 2, 0)
        axes[0, i].imshow(np.clip(orig, 0, 1))
        axes[0, i].axis('off')
        if i == 0:
            axes[0, i].set_title('Original', fontsize=12)
        
        # Reconstruction
        rec = recon[i].cpu().numpy().transpose(1, 2, 0)
        axes[1, i].imshow(np.clip(rec, 0, 1))
        axes[1, i].axis('off')
        if i == 0:
            axes[1, i].set_title('Reconstructed', fontsize=12)
    
    plt.suptitle('Original vs Reconstructed Images', fontsize=14, fontweight='bold')
    plt.tight_layout()
    plt.savefig('reconstructions.png', dpi=150, bbox_inches='tight')
    plt.show()
    
    print("Reconstruction comparison saved to 'reconstructions.png'")

# Show reconstructions
print("Comparing original images with their reconstructions...")
show_reconstructions(model, test_loader, num_images=8)

Comparing original images with their reconstructions...

Reconstruction comparison saved to 'reconstructions.png'

def latent_space_interpolation(model, num_steps=10, num_rows=3):
    """
    Perform linear interpolation between random latent vectors.
    
    Creates multiple rows of interpolations to show different transitions.
    Each step is clearly labeled for evaluation.
    """
    model.eval()
    
    fig, axes = plt.subplots(num_rows, num_steps, figsize=(num_steps*1.5, num_rows*1.5))
    
    for row in range(num_rows):
        # Sample two random latent vectors
        z1 = torch.randn(1, model.latent_dim).to(device)
        z2 = torch.randn(1, model.latent_dim).to(device)
        
        # Create interpolation steps (linear interpolation)
        alphas = torch.linspace(0, 1, num_steps)
        
        for i, alpha in enumerate(alphas):
            # Linear interpolation: z = (1-α)*z1 + α*z2
            z_interp = (1 - alpha) * z1 + alpha * z2
            
            # Generate image
            with torch.no_grad():
                img = model.decoder(z_interp)
            
            # Denormalize and display
            img = (img * 0.5 + 0.5).cpu().squeeze(0).numpy().transpose(1, 2, 0)
            img = np.clip(img, 0, 1)
            
            axes[row, i].imshow(img)
            axes[row, i].axis('off')
            
            # Add clear step labels on first row
            if row == 0:
                axes[row, i].set_title(f'Step {i+1}/{num_steps}', fontsize=9, fontweight='bold')
            
            # Add alpha values on last row
            if row == num_rows - 1:
                axes[row, i].set_xlabel(f'α={alpha:.2f}', fontsize=8)
    
    plt.suptitle(f'Latent Space Interpolation ({num_steps} Steps)', 
                 fontsize=14, fontweight='bold')
    plt.tight_layout()
    plt.savefig('interpolation.png', dpi=150, bbox_inches='tight')
    plt.show()
    
    print(f"Interpolation visualization saved showing {num_steps} clear steps")

# Perform interpolation
print("Exploring the latent space through interpolation...")
latent_space_interpolation(model, num_steps=10, num_rows=3)

Exploring the latent space through interpolation...

Interpolation visualization saved showing 10 clear steps

model_beta1 = model
history_beta1 = history

print("=" * 60)
print("TRAINING β-VAE (β = 5)")
print("=" * 60)

model_beta5 = VAE(LATENT_DIM).to(device)
history_beta5 = train_vae(model_beta5, train_loader, EPOCHS, beta=5.0, lr=LEARNING_RATE)

print("\nβ-VAE training complete!")

============================================================
TRAINING β-VAE (β = 5)
============================================================

Epoch 1/50: 100%|██████████| 391/391 [00:12<00:00, 30.76it/s, loss=360.6312, recon=294.3878, kl=13.2487]

  → Avg Loss: 498.0148 | Recon: 438.0461 | KL: 11.9937

Epoch 2/50: 100%|██████████| 391/391 [00:13<00:00, 30.07it/s, loss=347.5817, recon=270.5161, kl=15.4131]

  → Avg Loss: 367.7309 | Recon: 292.5369 | KL: 15.0388

Epoch 3/50: 100%|██████████| 391/391 [00:12<00:00, 30.14it/s, loss=319.2070, recon=244.1280, kl=15.0158]

  → Avg Loss: 349.2755 | Recon: 269.6205 | KL: 15.9310

Epoch 4/50: 100%|██████████| 391/391 [00:12<00:00, 30.39it/s, loss=346.5734, recon=270.2445, kl=15.2658]

  → Avg Loss: 342.5169 | Recon: 260.8127 | KL: 16.3408

Epoch 5/50: 100%|██████████| 391/391 [00:12<00:00, 30.31it/s, loss=338.4818, recon=253.3309, kl=17.0302]

def plot_beta_comparison(history_beta1, history_beta5):
    """Create side-by-side comparison plots for β=1 vs β=5 training"""
    fig, axes = plt.subplots(1, 3, figsize=(15, 4))
    epochs = range(1, len(history_beta1['total_loss']) + 1)
    
    # Total Loss
    axes[0].plot(epochs, history_beta1['total_loss'], 'b-', label='β=1', linewidth=2)
    axes[0].plot(epochs, history_beta5['total_loss'], 'r-', label='β=5', linewidth=2)
    axes[0].set_xlabel('Epoch', fontsize=11)
    axes[0].set_ylabel('Loss', fontsize=11)
    axes[0].set_title('Total Loss Comparison', fontsize=12, fontweight='bold')
    axes[0].legend(fontsize=10)
    axes[0].grid(True, alpha=0.3)
    
    # Reconstruction Loss
    axes[1].plot(epochs, history_beta1['recon_loss'], 'b-', label='β=1', linewidth=2)
    axes[1].plot(epochs, history_beta5['recon_loss'], 'r-', label='β=5', linewidth=2)
    axes[1].set_xlabel('Epoch', fontsize=11)
    axes[1].set_ylabel('Loss', fontsize=11)
    axes[1].set_title('Reconstruction Loss Comparison', fontsize=12, fontweight='bold')
    axes[1].legend(fontsize=10)
    axes[1].grid(True, alpha=0.3)
    
    # KL Divergence
    axes[2].plot(epochs, history_beta1['kl_loss'], 'b-', label='β=1', linewidth=2)
    axes[2].plot(epochs, history_beta5['kl_loss'], 'r-', label='β=5', linewidth=2)
    axes[2].set_xlabel('Epoch', fontsize=11)
    axes[2].set_ylabel('Loss', fontsize=11)
    axes[2].set_title('KL Divergence Comparison', fontsize=12, fontweight='bold')
    axes[2].legend(fontsize=10)
    axes[2].grid(True, alpha=0.3)
    
    plt.suptitle('β-VAE Training Comparison: Direct Overlay', fontsize=14, fontweight='bold')
    plt.tight_layout()
    plt.savefig('beta_comparison_overlay.png', dpi=150, bbox_inches='tight')
    plt.show()
    
    print("β comparison overlay saved to 'beta_comparison_overlay.png'")

print("β comparison plotting function defined")

β comparison plotting function defined

# Plot individual β=5 training curves
print("Plotting β=5 training curves...")
plot_training_curves(history_beta5, "β-VAE (β=5) Training Progress")

Plotting β=5 training curves...

Training curves saved to 'training_curves.png'

# Plot overlay comparison of both models
print("Creating overlay comparison plot...")
plot_beta_comparison(history_beta1, history_beta5)

Creating overlay comparison plot...

β comparison overlay saved to 'beta_comparison_overlay.png'

# Compare training curves
fig, axes = plt.subplots(2, 3, figsize=(15, 8))

epochs = range(1, len(history_beta1['total_loss']) + 1)

# Row 1: β = 1
axes[0, 0].plot(epochs, history_beta1['total_loss'], 'b-', linewidth=2)
axes[0, 0].set_title('Total Loss (β=1)')
axes[0, 0].set_xlabel('Epoch')
axes[0, 0].grid(True, alpha=0.3)

axes[0, 1].plot(epochs, history_beta1['recon_loss'], 'g-', linewidth=2)
axes[0, 1].set_title('Reconstruction Loss (β=1)')
axes[0, 1].set_xlabel('Epoch')
axes[0, 1].grid(True, alpha=0.3)

axes[0, 2].plot(epochs, history_beta1['kl_loss'], 'r-', linewidth=2)
axes[0, 2].set_title('KL Loss (β=1)')
axes[0, 2].set_xlabel('Epoch')
axes[0, 2].grid(True, alpha=0.3)

# Row 2: β = 5
axes[1, 0].plot(epochs, history_beta5['total_loss'], 'b-', linewidth=2)
axes[1, 0].set_title('Total Loss (β=5)')
axes[1, 0].set_xlabel('Epoch')
axes[1, 0].grid(True, alpha=0.3)

axes[1, 1].plot(epochs, history_beta5['recon_loss'], 'g-', linewidth=2)
axes[1, 1].set_title('Reconstruction Loss (β=5)')
axes[1, 1].set_xlabel('Epoch')
axes[1, 1].grid(True, alpha=0.3)

axes[1, 2].plot(epochs, history_beta5['kl_loss'], 'r-', linewidth=2)
axes[1, 2].set_title('KL Loss (β=5)')
axes[1, 2].set_xlabel('Epoch')
axes[1, 2].grid(True, alpha=0.3)

plt.suptitle('Training Comparison: Standard VAE vs β-VAE', fontsize=14, fontweight='bold')
plt.tight_layout()
plt.savefig('beta_comparison_curves.png', dpi=150, bbox_inches='tight')
plt.show()

print("Comparison curves saved to 'beta_comparison_curves.png'")

Comparison curves saved to 'beta_comparison_curves.png'

# Compare generated images side by side
def compare_generations(model1, model2, title1="β=1", title2="β=5", num_images=8):
    """Generate images from both models for comparison"""
    model1.eval()
    model2.eval()
    
    # Use the SAME random seed for fair comparison
    torch.manual_seed(42)
    z = torch.randn(num_images, LATENT_DIM).to(device)
    
    with torch.no_grad():
        gen1 = model1.decoder(z)
        gen2 = model2.decoder(z)
    
    # Denormalize
    gen1 = (gen1 * 0.5 + 0.5).cpu()
    gen2 = (gen2 * 0.5 + 0.5).cpu()
    
    # Plot
    fig, axes = plt.subplots(2, num_images, figsize=(14, 4))
    
    for i in range(num_images):
        img1 = gen1[i].numpy().transpose(1, 2, 0)
        img2 = gen2[i].numpy().transpose(1, 2, 0)
        
        axes[0, i].imshow(np.clip(img1, 0, 1))
        axes[0, i].axis('off')
        if i == 0:
            axes[0, i].set_ylabel(title1, fontsize=12, rotation=0, labelpad=30)
        
        axes[1, i].imshow(np.clip(img2, 0, 1))
        axes[1, i].axis('off')
        if i == 0:
            axes[1, i].set_ylabel(title2, fontsize=12, rotation=0, labelpad=30)
    
    plt.suptitle('Generated Image Comparison (Same Latent Vectors)', fontsize=14, fontweight='bold')
    plt.tight_layout()
    plt.savefig('beta_comparison_images.png', dpi=150, bbox_inches='tight')
    plt.show()
    
    print("Image comparison saved to 'beta_comparison_images.png'")

# Compare!
print("Comparing generated images from both models...")
compare_generations(model_beta1, model_beta5)

Comparing generated images from both models...

Image comparison saved to 'beta_comparison_images.png'

# Compare reconstruction quality
def compare_reconstructions(model1, model2, data_loader, title1="β=1", title2="β=5", num_images=6):
    """Compare reconstruction quality of both models"""
    model1.eval()
    model2.eval()
    
    images, _ = next(iter(data_loader))
    images = images[:num_images].to(device)
    
    with torch.no_grad():
        recon1, _, _ = model1(images)
        recon2, _, _ = model2(images)
    
    # Denormalize
    images = (images * 0.5 + 0.5).cpu()
    recon1 = (recon1 * 0.5 + 0.5).cpu()
    recon2 = (recon2 * 0.5 + 0.5).cpu()
    
    # Plot
    fig, axes = plt.subplots(3, num_images, figsize=(12, 6))
    
    for i in range(num_images):
        orig = images[i].numpy().transpose(1, 2, 0)
        r1 = recon1[i].numpy().transpose(1, 2, 0)
        r2 = recon2[i].numpy().transpose(1, 2, 0)
        
        axes[0, i].imshow(np.clip(orig, 0, 1))
        axes[0, i].axis('off')
        if i == 0:
            axes[0, i].set_ylabel('Original', fontsize=10, rotation=0, labelpad=40)
        
        axes[1, i].imshow(np.clip(r1, 0, 1))
        axes[1, i].axis('off')
        if i == 0:
            axes[1, i].set_ylabel(f'Recon\n({title1})', fontsize=10, rotation=0, labelpad=40)
        
        axes[2, i].imshow(np.clip(r2, 0, 1))
        axes[2, i].axis('off')
        if i == 0:
            axes[2, i].set_ylabel(f'Recon\n({title2})', fontsize=10, rotation=0, labelpad=40)
    
    plt.suptitle('Reconstruction Quality Comparison', fontsize=14, fontweight='bold')
    plt.tight_layout()
    plt.savefig('beta_comparison_recon.png', dpi=150, bbox_inches='tight')
    plt.show()
    
    print("Reconstruction comparison saved to 'beta_comparison_recon.png'")

compare_reconstructions(model_beta1, model_beta5, test_loader)

Reconstruction comparison saved to 'beta_comparison_recon.png'

print("Generating 16-image grid for β-VAE (β=5)...")
generate_image_grid(model_beta5, num_images=16, nrow=4, title="β-VAE Generated Images (β=5)")

Generating 16-image grid for β-VAE (β=5)...

Generated image grid saved to 'generated_grid.png'

print("=" * 60)
print("FINAL TRAINING STATISTICS")
print("=" * 60)

print("\nStandard VAE (β=1):")
print(f"   Final Total Loss: {history_beta1['total_loss'][-1]:.4f}")
print(f"   Final Recon Loss: {history_beta1['recon_loss'][-1]:.4f}")
print(f"   Final KL Loss: {history_beta1['kl_loss'][-1]:.4f}")

print("\nβ-VAE (β=5):")
print(f"   Final Total Loss: {history_beta5['total_loss'][-1]:.4f}")
print(f"   Final Recon Loss: {history_beta5['recon_loss'][-1]:.4f}")
print(f"   Final KL Loss: {history_beta5['kl_loss'][-1]:.4f}")

print("\n" + "=" * 60)
print("Assignment Complete!")
print("=" * 60)

============================================================
FINAL TRAINING STATISTICS
============================================================

Standard VAE (β=1):
   Final Total Loss: 185.4197
   Final Recon Loss: 125.6657
   Final KL Loss: 59.7539

β-VAE (β=5):
   Final Total Loss: 315.4138
   Final Recon Loss: 222.5905
   Final KL Loss: 18.5647

============================================================
Assignment Complete!
============================================================

# Save individual model weights locally
torch.save(model_beta1.state_dict(), 'vae_beta1_weights.pth')
torch.save(model_beta5.state_dict(), 'vae_beta5_weights.pth')

print("Saved weights to 'vae_beta1_weights.pth' and 'vae_beta5_weights.pth'")

Saved weights to 'vae_beta1_weights.pth' and 'vae_beta5_weights.pth'

torch.save({
    'model_beta1_state_dict': model_beta1.state_dict(),
    'model_beta5_state_dict': model_beta5.state_dict(),
    'history_beta1': history_beta1,
    'history_beta5': history_beta5,
    'latent_dim': LATENT_DIM,
    'epochs': EPOCHS,
}, 'vae_models.pth')

print("Models saved to 'vae_models.pth'")

Models saved to 'vae_models.pth'

from PIL import Image
import os

def display_saved_image(filepath, title):
    """Display a saved image if it exists"""
    if os.path.exists(filepath):
        img = Image.open(filepath)
        plt.figure(figsize=(15, 10))
        plt.imshow(img)
        plt.axis('off')
        plt.title(title, fontsize=16, fontweight='bold', pad=20)
        plt.tight_layout()
        plt.show()
        print(f"✓ Displayed: {filepath}\n")
    else:
        print(f"✗ File not found: {filepath}\n")

print("=" * 60)
print("DISPLAYING ALL SAVED OUTPUTS")
print("=" * 60)
print()

============================================================
DISPLAYING ALL SAVED OUTPUTS
============================================================

display_saved_image('training_curves.png', 'Training Progress: Standard VAE (β=1)')

✓ Displayed: training_curves.png

display_saved_image('generated_grid.png', 'VAE Generated Images (Random Noise → Images)')

✓ Displayed: generated_grid.png

display_saved_image('reconstructions.png', 'Original vs Reconstructed Images')

✓ Displayed: reconstructions.png

display_saved_image('interpolation.png', 'Latent Space Interpolation: Smooth Morphing Between Images')

✓ Displayed: interpolation.png

display_saved_image('beta_comparison_curves.png', 'β-VAE Training Comparison: β=1 vs β=5')

✓ Displayed: beta_comparison_curves.png

display_saved_image('beta_comparison_images.png', 'Generated Images: Same Latent Vectors, Different β')

✓ Displayed: beta_comparison_images.png

display_saved_image('beta_comparison_recon.png', 'Reconstruction Quality: β=1 vs β=5')

✓ Displayed: beta_comparison_recon.png

Variational Autoencoder (VAE) for CIFAR-10 Color Image Synthesis¶

Foundational Models and Generative AI - Assignment 1¶

I have implemented below components¶

This image explains what exactly I have done¶

I have created streamlit app for VAE simulation¶

1.Import Libraries¶

2. Load CIFAR-10 Dataset¶

3. Build the Encoder Network¶

4. Decoder Network¶

5. The Complete VAE with Reparameterization Trick¶

The Problem:¶

The Solution (Reparameterization Trick):¶

6. Loss Function (ELBO)¶

1. Reconstruction Loss¶

2. KL Divergence Loss¶

Total Loss (ELBO)¶

7. Training Loop¶

8. Training Curves Visualization¶

9. Generate Images from Random Noise (16-Image Grid)¶

10. Reconstruction Quality Check¶

11. Latent Space Interpolation (Task 3)¶

12. β-VAE Experiment (Task 4)¶

13. Comparing β=1 vs β=5¶

Observations on β-VAE (β=1 vs β=5)¶

15. Final Summary & Saved Outputs¶

Submission¶

16. Display All Generated Images¶

1. Training Curves (Standard VAE β=1)¶

2. Generated Images (16-Image Grid)¶

3. Image Reconstructions¶

4. Latent Space Interpolation (10 Steps)¶

5. β-VAE Comparison: Training Curves¶

6. β-VAE Comparison: Generated Images¶

7. β-VAE Comparison: Reconstruction Quality¶